#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：gas_pipeline_qc_partition3.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/23 16:17 
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
QC-only Scheduling with Two Analyzers per Component
---------------------------------------------------
- 不考虑充气分次：同配方视为已一次性完成充装，所有瓶 mix_end=0，直接进入分析。
- 分析(QC)：
  * 每种组分有两台设备（如 CH4-1/CH4-2），每台只测这一种组分；
  * 每瓶的组分可被划分为若干“检测组”：组大小=1 → 串行；组大小>=2 → 1-pass（同起同止），
    段长=该组内各设备的单机时长的最大值，中途不允许拆换。
- GA 优化每瓶的分组策略（0=ALL, 1=SPLIT_LONGEST, 2=EACH），目标最小化 Cmax。
- 输出：收敛曲线 + 甘特图（文字黑色、自动换行）。

Author: you
"""
import random
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple, Optional

try:
    import matplotlib.pyplot as plt
    HAS_MPL = True
except Exception:
    HAS_MPL = False


# ============================ CONFIG（可自定义） ============================ #
CONFIG = {
    # 规模
    "NUM_BOTTLES": 20,                 # 气瓶数量（工件数）
    "MAX_COMPONENTS_PER_BOTTLE": 4,    # 每瓶最大组分数（1..K）

    # 组分全集（自动为每个组分生成两台设备）
    "COMPONENT_TYPES": ["A","B","C","D","O2","CO2","CH4","H2","CO","N2"],

    # QC 单机检测时间（小时）
    "COMPONENT_TEST_TIME": {
        "A":0.22, "B":0.20, "C":0.25, "D":0.18,
        "O2":0.18, "CO2":0.22, "CH4":0.25, "H2":0.15, "CO":0.20, "N2":0.12
    },

    # QC 换型时间（同一设备上，如果上一次签名≠这次的组件名，则需要 setup）
    "QC_SETUP_TIME": 0.12,

    # GA 参数（优化每瓶分组策略）
    "GA_POP": 80,
    "GA_GENS": 200,
    "GA_PC": 0.9,
    "GA_PM": 0.35,
    "GA_ELITE": 10,

    # 绘图
    "LABEL_WRAP_CHARS": 22,
    "RANDOM_SEED": 7
}

# ============================ 数据结构 ============================ #
@dataclass(frozen=True)
class Recipe:
    components: Tuple[str, ...]   # 不考虑浓度，只要组分相同视作同配方

    def key(self):
        return tuple(sorted(self.components))

@dataclass
class Bottle:
    id: int
    recipe: Recipe

@dataclass
class QCTask:
    start: float
    end: float
    device: int           # 设备全局索引
    label: str
    group_key: Optional[Tuple[int,float,float,int]] = None  # (bid, st, en, group_idx) for 1-pass coloring

# ============================ 随机实例生成 ============================ #
class RandomData:
    """
    - 为每个组分生成两台设备（component -> device ids）
    - 为每瓶从组件全集中抽 1..K 个组分
    """
    def __init__(self, cfg):
        self.cfg = cfg
        random.seed(cfg["RANDOM_SEED"]); np.random.seed(cfg["RANDOM_SEED"])

    def gen(self):
        comps_all = self.cfg["COMPONENT_TYPES"]
        # 创建设备：每个组分两台
        comp_to_devices: Dict[str, List[int]] = {}
        device_labels: List[str] = []
        device_index = 0
        for c in comps_all:
            comp_to_devices[c] = [device_index, device_index+1]
            device_labels.append(f"{c}-1")
            device_labels.append(f"{c}-2")
            device_index += 2
        D = device_index  # 总设备数 = 2 * |components|

        # 生成气瓶
        N = self.cfg["NUM_BOTTLES"]
        Kmax = self.cfg["MAX_COMPONENTS_PER_BOTTLE"]
        bottles: List[Bottle] = []
        for i in range(N):
            k = np.random.randint(1, Kmax+1)
            k = min(k, len(comps_all))
            comps = random.sample(comps_all, k)
            bottles.append(Bottle(i, Recipe(tuple(comps))))
        return bottles, comp_to_devices, device_labels

# ============================ 调度器（仅 QC） ============================ #
class Scheduler:
    """
    仅分析阶段调度：
    - 每瓶的组分按“有序分组列表”执行：
      * 组大小=1：串行 → 选该组分的两台设备中“就绪+setup 最早”的那台
      * 组大小>=2：1-pass → 对每个组分各选其两台设备中“就绪+setup 最早”的那台；所有设备同起同止
        段长=组内单机时长的最大值，中途不允许拆换
    """
    def __init__(self, cfg, bottles: List[Bottle], comp_to_devices: Dict[str, List[int]], device_labels: List[str]):
        self.cfg = cfg
        self.bottles = bottles
        self.N = len(bottles)
        self.comp_to_devices = comp_to_devices
        self.device_labels = device_labels
        self.D = len(device_labels)

    def single_test_time(self, comp: str) -> float:
        return float(self.cfg["COMPONENT_TEST_TIME"].get(comp, 0.2))

    # policy codes: 0=ALL_TOGETHER, 1=SPLIT_LONGEST, 2=EACH_SEPARATE
    def make_groups(self, comps: List[str], policy: int) -> List[List[str]]:
        if len(comps) <= 1:
            return [comps[:]]
        t = {c: self.single_test_time(c) for c in comps}
        if policy == 0:   # ALL together
            return [comps[:]]
        if policy == 1:   # longest alone + rest together
            longest = max(comps, key=lambda x: t[x])
            rest = [c for c in comps if c != longest]
            return [[longest], rest] if rest else [[longest]]
        # EACH separate
        return [[c] for c in sorted(comps)]

    def schedule_qc_with_groups(self, bottle_groups: Dict[int, List[List[str]]]) -> Tuple[float, List[QCTask]]:
        D = self.D
        ready = [0.0]*D
        last_sig: List[Optional[str]] = [None]*D    # 上次在该设备测的组分名
        setup_t = self.cfg["QC_SETUP_TIME"]

        tasks: List[QCTask] = []
        # 本题不考虑充气排队，默认充气同时完成 → 所有瓶 mix_end=0
        # 可以按瓶 id 顺序或随机进入 QC；这里按 id 顺序
        order_seq = sorted(self.bottles, key=lambda b: b.id)

        for b in order_seq:
            bid = b.id
            t_cur = 0.0
            groups = bottle_groups[bid]  # 有序分组

            for gi, group in enumerate(groups):
                if len(group) == 1:
                    comp = group[0]
                    devs = self.comp_to_devices.get(comp, [])
                    if not devs:
                        raise RuntimeError(f"No device for component {comp}")
                    best_d=None; best_ready=None
                    for d in devs:
                        stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                        rdy = ready[d] + stp
                        if (best_ready is None) or (rdy < best_ready):
                            best_d, best_ready = d, rdy
                    st = max(t_cur, best_ready)
                    pt = self.single_test_time(comp)
                    en = st + pt
                    label = f"S{bid} serial {self.device_labels[best_d]}({comp})"
                    tasks.append(QCTask(st, en, best_d, label, None))
                    ready[best_d] = en
                    last_sig[best_d] = comp
                    t_cur = en
                    continue

                # 1-pass 组：每个组分挑一台该组分的两台设备中“就绪+setup 最早”的那台
                chosen: List[Tuple[int, str, float]] = []  # (device, comp, ready_with_setup)
                earliest = t_cur
                for comp in group:
                    devs = self.comp_to_devices.get(comp, [])
                    if not devs:
                        raise RuntimeError(f"No device for component {comp}")
                    best_d=None; best_ready=None
                    for d in devs:
                        stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                        rdy = ready[d] + stp
                        if (best_ready is None) or (rdy < best_ready):
                            best_d, best_ready = d, rdy
                    chosen.append((best_d, comp, best_ready))
                    earliest = max(earliest, best_ready)

                # 同起同止，段长取该组各组件单机时长的最大值（中途不允许换/拆）
                seg_len = max(self.single_test_time(c) for (_, c, _) in chosen)
                st = earliest; en = st + seg_len

                used = "+".join(self.device_labels[d] for (d,_,_) in chosen)
                comps_str = ",".join(sorted(group))
                label = f"S{bid} 1-pass[{used}] comps={comps_str} max={seg_len:.2f}h"
                group_key = (bid, round(st,6), round(en,6), gi)
                for (d, comp, _) in chosen:
                    tasks.append(QCTask(st, en, d, label, group_key))
                    ready[d] = en
                    last_sig[d] = comp
                t_cur = en

        Cmax = max((t.end for t in tasks), default=0.0)
        return Cmax, tasks

# ============================ GA（优化每瓶分组策略） ============================ #
class GA:
    """
    染色体：长度 N（瓶数），每位∈{0,1,2}
      0=ALL（所有组分一次通过）
      1=SPLIT_LONGEST（最长单独，其余一起）
      2=EACH（每个组分分别检测）
    目标：最小化 Cmax
    """
    def __init__(self, cfg, scheduler: Scheduler):
        self.cfg = cfg
        self.S  = scheduler
        self.N  = self.S.N
        random.seed(cfg["RANDOM_SEED"]); np.random.seed(cfg["RANDOM_SEED"])

    def rand_ind(self):
        return np.random.randint(0,3,size=self.N,dtype=int)

    def crossover(self, a, b):
        n=len(a); cut=np.random.randint(1,n)
        c1=np.concatenate([a[:cut], b[cut:]])
        c2=np.concatenate([b[:cut], a[cut:]])
        return c1,c2

    def mutate(self, ind, pm=0.2):
        if np.random.rand()<pm:
            i=np.random.randint(self.N)
            ind[i]=np.random.randint(0,3)
        return ind

    def ind_to_groups(self, ind) -> Dict[int, List[List[str]]]:
        mp={}
        for i,b in enumerate(self.S.bottles):
            comps=list(b.recipe.components)
            mp[b.id]=self.S.make_groups(comps, int(ind[i]))
        return mp

    def fitness(self, ind):
        groups = self.ind_to_groups(ind)
        Cmax, _ = self.S.schedule_qc_with_groups(groups)
        return Cmax

    def run(self):
        pop=self.cfg["GA_POP"]; gens=self.cfg["GA_GENS"]
        pc=self.cfg["GA_PC"];   pm=self.cfg["GA_PM"]; elite=self.cfg["GA_ELITE"]

        pool=[]
        for _ in range(pop):
            ind=self.rand_ind()
            f=self.fitness(ind)
            pool.append((ind,f))
        best=min(pool, key=lambda x:x[1]); hist=[best[1]]

        for g in range(1, gens+1):
            pool.sort(key=lambda x:x[1])
            elites=pool[:elite]
            new=[(e[0].copy(), e[1]) for e in elites]
            while len(new)<pop:
                p1=random.choice(pool)[0]
                p2=random.choice(pool)[0]
                c1,c2=p1.copy(),p2.copy()
                if np.random.rand()<pc:
                    c1,c2=self.crossover(p1,p2)
                if np.random.rand()<pm:
                    c1=self.mutate(c1)
                if np.random.rand()<pm:
                    c2=self.mutate(c2)
                new.append((c1, self.fitness(c1)))
                if len(new)<pop:
                    new.append((c2, self.fitness(c2)))
            pool=new
            cur=min(pool, key=lambda x:x[1])
            if cur[1] < best[1]:
                best=cur
            hist.append(best[1])
        return best[0], best[1], hist

# ============================ 绘图 ============================ #
def wrap_text(s: str, width: int) -> str:
    out=[]; line=""
    for ch in s:
        line += ch
        if len(line) >= width and ch == ' ':
            out.append(line.rstrip()); line=""
    if line: out.append(line)
    return "\n".join(out)

def plot_convergence(hist, title="GA Convergence (minimize Cmax)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return
    plt.figure(figsize=(6,4))
    plt.plot(hist, linewidth=2, color='black')
    plt.xlabel("Generation", color='black'); plt.ylabel("Best Cmax", color='black')
    plt.title(title, color='black'); plt.grid(True, alpha=0.3, color='black')
    ax=plt.gca()
    ax.tick_params(axis='x', colors='black'); ax.tick_params(axis='y', colors='black')
    for spine in ax.spines.values(): spine.set_edgecolor('black')
    plt.tight_layout(); plt.show()

def plot_gantt(qc_tasks: List[QCTask], device_labels: List[str], cfg,
               title="QC Gantt (1-pass groups shown with same start/end)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return

    D = len(device_labels)
    row_idx = {d: d for d in range(D)}

    # 为1-pass组统一颜色
    from itertools import cycle
    palette = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
    grp_color: Dict[Tuple[int,float,float,int], str] = {}
    for t in qc_tasks:
        if t.group_key and t.group_key not in grp_color:
            grp_color[t.group_key] = next(palette)

    fig_h = max(4.5, 0.6*D + 2)
    fig, ax = plt.subplots(figsize=(12, fig_h))

    for t in qc_tasks:
        y=row_idx[t.device]
        color = grp_color.get(t.group_key, None)
        ax.barh(y, t.end-t.start, left=t.start, height=0.8,
                color=color, edgecolor='black', linewidth=0.6)
        ax.text((t.start+t.end)/2, y, wrap_text(t.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center', color='black', fontsize=8)

    ax.set_yticks(list(range(D))); ax.set_yticklabels([f"E{d+1} ({device_labels[d]})" for d in range(D)], color='black')
    ax.set_xlabel("Time (h)", color='black'); ax.set_title(title, color='black')
    ax.tick_params(axis='x', colors='black'); ax.tick_params(axis='y', colors='black')
    ax.grid(axis='x', alpha=0.3, color='black')
    for spine in ax.spines.values(): spine.set_edgecolor('black')
    plt.tight_layout(); plt.show()

# ============================ 主程序 ============================ #
def main():
    random.seed(CONFIG["RANDOM_SEED"]); np.random.seed(CONFIG["RANDOM_SEED"])

    # 1) 实例
    bottles, comp_to_devices, device_labels = RandomData(CONFIG).gen()

    # 2) 调度器（仅QC）
    sch = Scheduler(CONFIG, bottles, comp_to_devices, device_labels)

    # 3) GA 优化分组策略
    ga = GA(CONFIG, sch)
    best_ind, best_cmax, hist = ga.run()

    # 4) 用最优策略生成分组并排程
    bottle_groups={}
    for i,b in enumerate(bottles):
        comps=list(b.recipe.components)
        bottle_groups[b.id] = sch.make_groups(comps, int(best_ind[i]))
    Cmax, qc_tasks = sch.schedule_qc_with_groups(bottle_groups)

    print(f"Bottles={CONFIG['NUM_BOTTLES']}, QC devices={len(device_labels)} (2 per component)")
    print("Best policy per bottle (0=ALL,1=SPLIT_LONGEST,2=EACH):")
    print(best_ind.tolist())
    print(f"Best Cmax: {Cmax:.3f} h")

    if HAS_MPL:
        plot_convergence(hist)
        plot_gantt(qc_tasks, device_labels, CONFIG)
    else:
        print("(Install matplotlib to view plots.)")

if __name__ == "__main__":
    main()
